import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
import cv2
from insightface.app import FaceAnalysis
from pymilvus import MilvusClient, DataType

cli = MilvusClient(
    uri='http://39.104.78.210:19530',
    token='root:Milvus',
) # 建立连接
print(cli)

cli.drop_collection(collection_name='faceD') # 删除数据库
schema = cli.create_schema(
    auto_id=False,
    enable_dynamic_field=True
) # 创建的表结构, 包括如下操作

# schema.add_field(field_name='id', datatype=DataType.INT64, is_primary=True, auto_id=True)
schema.add_field(field_name='key', datatype=DataType.VARCHAR, is_primary=True, max_length=30)
schema.add_field(field_name='vector', datatype=DataType.FLOAT_VECTOR, dim=512)
index_params = cli.prepare_index_params() # 创建的索引结构 包括如下操作

index_params.add_index(
    field_name='key',
    index_type='AUTOINDEX'
)
index_params.add_index(
    field_name='vector',
    index_type='AUTOINDEX',
    metric_type="COSINE"
)
cli.create_collection(
    collection_name='faceD',
    schema=schema,
    index_params=index_params
) # 创建collection
#
# cli.create_collection(
#     collection_name='demo',
#     dimension=2,
# )
# vectors = []
# # imgs = []
# # vectors.append([1, 1])
# # for i in range(200, 300):
# #     vectors.append([i, i])
# # # 填入 1，1  200，200 201，201 。。。
# app = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])  # 优先GPU加速
# app.prepare(ctx_id=0, det_size=(640, 640))  # 第一个使用GPU， 检测分辨率
#
# for i in range(1, 6):
#     img = cv2.imread(str(i) + '.jpg')
#     faces = app.get(img)
#     for j, face in enumerate(faces):
#         # print(i, face)
#         embedding = face.embedding
#         vectors.append(embedding)
#         # print(embedding.shape)
#
# data = [
#     {'name': "阚思涵", 'vector': vectors[0]},
#     {'name': "刘世鹏", 'vector': vectors[1]},
#     {'name': "乔雨琪", 'vector': vectors[2]},
#     {'name': "丁海乐", 'vector': vectors[3]},
#     {'name': "王恩意", 'vector': vectors[4]},
# ]
# #
# # res = cli.insert(
# #     collection_name="faceD",
# #     data=data,
# # )
# # print(res)
#
# res = cli.search(
#     collection_name="faceD",
#     anns_field="vector",
#     data=[vectors[0]],
#     limit=1,
#     output_fields=["name"],
#     search_params={"metric_type": "COSINE"}
# ) # ANN(KNN) 查找最相近的元素
# print(res)
#
